import aio_pika
from functools import partial

from .base import BaseBroker
from pubsub.utils import ensure_bytes, random_string


class RabbitConnection:
    def __init__(self, amqp_url, **connection_kwargs):
        self.amqp_url = amqp_url
        self.connection = None
        self.connection_kwargs = connection_kwargs

    async def start(self):
        self.connection = await aio_pika.connect_robust(
            self.amqp_url, **self.connection_kwargs)

    async def stop(self):
        await self.connection.stop()
        self.connection = None

    async def channel(self):
        return await self.connection.channel()


class RabbitMQ(BaseBroker):
    def __init__(self, *args, exchange_name=None, **kwargs):
        super().__init__(*args, **kwargs)
        assert exchange_name is not None
        self.exchange_name = exchange_name
        self.consumer_tag = self.tag + '_' + random_string()
        self.queue_name = f'{self.exchange_name}_{self.consumer_tag}'
        self.log_context['consumer_tag'] = self.consumer_tag

        self.channel = None
        self.exchange = None
        self.queue = None

    @property
    def is_closed(self):
        if not self.channel:
            return True
        return self.channel.is_closed 

    async def declare_exchagne(self):
        if self.channel is None:
            self.channel = await self.connection.channel()
            self.logger.debug('channel declared')
        if self.exchange is None:
            self.exchagne = await self.channel.declare_exchange(
                name=self.exchange_name, type=aio_pika.ExchangeType.HEADERS)
            self.logger.debug('exchange declared')

    async def start_publish(self):
        await self.declare_exchagne()

    async def publish(self, message, *channel_args, **channel_kwargs):
        message = ensure_bytes(message)
        for key, value in channel_kwargs.items():
            headers = {key: value}
            pika_message = aio_pika.Message(
                message,
                delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
                headers=headers)
            await self.exchagne.publish(pika_message, routing_key='')
            self.logger.debug('publish headers: %s, message: %s', headers,
                              message)

    async def start_consume(self, channels_subscribed, callback):
        await self.declare_exchagne()
        assert self.queue is None, 'start_consume already called'
        self.queue = await self.channel.declare_queue(
            self.queue_name,
            exclusive=True,
            auto_delete=True,
        )
        self.logger.debug(f'queue {self.queue_name} declared')
        channels_subscribed['x-match'] = 'any'
        await self.queue.bind(self.exchagne, arguments=channels_subscribed)
        message_callback = partial(self.consume_message, callback=callback)
        await self.queue.consume(
            message_callback,
            consumer_tag=self.consumer_tag,
        )
        self.logger.debug('ready to consume')

    async def consume_message(self, message, callback=None):
        async with message.process():
            decoded_body = message.body.decode()
            self.logger.debug(f'[consume_message] {decoded_body}')
            if callback is None:
                self.logger.warning('[consume_message] callback is None')
                return
        try:
            await callback(decoded_body)
        except Exception as err:
            self.logger.exception('[consume_message] unexpected err: %s',
                                  err,
                                  exc_info=True)

    async def stop(self):
        if self.channel:
            self.logger.info('channel is closed')
            await self.channel.close()
        self.logger.debug('closed')
